from datetime import datetime


def calculate_pyramid_add(total_amount, initial_price, num_adds, drop_percent, ratio=1.1):
    """
    金字塔加仓计算器
    :param total_amount: 可加仓总金额
    :param initial_price: 初始加仓价位
    :param num_adds: 加仓次数
    :param drop_percent: 每次下跌百分比（如5表示下跌5%）
    :param ratio: 加仓系数比，默认1.1倍
    :return: 按价位排序的加仓计划列表 [(price, amount, shares), ...]
    """
    # 计算加仓权重（等比数列）
    weights = [ratio ** i for i in range(num_adds)]
    total_weight = sum(weights)

    # 计算基础加仓金额
    base_amount = total_amount / total_weight

    # 生成加仓计划
    add_plan = []
    current_price = initial_price
    for i in range(num_adds):
        if i > 0:
            current_price *= (1 - drop_percent / 100)

        # 计算理论加仓金额
        theoretical_amount = base_amount * weights[i]

        # 计算实际购买股数（100股取整）
        shares = int(theoretical_amount / current_price)
        shares = max(round(shares / 100) * 100, 100)  # 至少100股

        # 计算实际加仓金额
        actual_amount = round(current_price * shares, 2)

        add_plan.append((
            round(current_price, 3),
            actual_amount,
            shares
        ))

    # 按价格降序排序
    return sorted(add_plan, key=lambda x: -x[0])


if __name__ == '__main__':
    total_amount = float(input("请输入可加仓总金额: "))
    initial_price = float(input("请输入初始加仓价位: "))
    num_adds = int(input("请输入加仓次数: "))
    drop_percent = float(input("请输入下跌百分比（如输入5表示下跌5%）: "))
    ratio = float(input("请输入加仓系数比（默认1.1倍）: ") or 1.1)
    #excel文件名
    excel_file_name = input("请输入excel文件名（默认加仓计划.xlsx）: ") or "加仓计划.xlsx"

    # 确保文件名有.xlsx扩展名
    if not excel_file_name.endswith('.xlsx'):
        excel_file_name += '.xlsx'


    plan = calculate_pyramid_add(
        total_amount=total_amount,
        initial_price=initial_price,
        num_adds=num_adds,
        drop_percent=drop_percent,
        ratio=ratio
    )

    print("\n加仓计划：")
    print("价位\t\t加仓金额\t\t股数")
    total_actual = 0
    for price, amount, shares in plan:
        print(f"{price:.3f}\t\t{amount:.2f}\t\t{shares}")
        total_actual += amount

    #最后结果输出excel文件
    import pandas as pd
    # 创建输出目录
    import os

    df = pd.DataFrame(plan, columns=['价位', '加仓金额', '股数'])
    dateStamp = datetime.now().strftime("%Y%m%d")
    screenshot_dir = f"../spiltExcel/output/加仓计划"
    os.makedirs(screenshot_dir, exist_ok=True)


    df.to_excel(f"{screenshot_dir}/{excel_file_name}", index=False)
    print(f"\n已将结果保存到 {screenshot_dir}/{excel_file_name} 文件")

    print(f"\n实际总加仓金额: {total_actual:.2f}")
    print(f"理论剩余金额: {total_amount - total_actual:.2f}")
    print("注：实际金额可能因股数取整与理论值存在差异")